{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6.5 循环神经网络的简洁实现\n",
    "\n",
    "本节将使用Keras来更简洁地实现基于循环神经网络的语言模型。首先，我们读取周杰伦专辑歌词数据集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "1"
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import backend as f\n",
    "import numpy as np\n",
    "import math\n",
    "import sys\n",
    "import time\n",
    "sys.path.append(\"..\") \n",
    "import d2lzh_tensorflow2 as d2l\n",
    "\n",
    "(corpus_indices, char_to_idx, idx_to_char, vocab_size) = d2l.load_data_jay_lyrics()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义模型\n",
    "\n",
    "Keras的`Rnn`模块提供了循环神经网络的实现。下面构造一个含单隐藏层、隐藏单元个数为256的循环神经网络层`rnn_layer`，并对权重做初始化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "26"
    }
   },
   "outputs": [],
   "source": [
    "num_hiddens = 256\n",
    "cell=keras.layers.SimpleRNNCell(num_hiddens,kernel_initializer='glorot_uniform')\n",
    "rnn_layer = keras.layers.RNN(cell,time_major=True,return_sequences=True,return_state=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来调用`rnn_layer`的成员函数`get_initial_state`来返回初始化的隐藏状态列表。它有一个形状为( 批量大小, 隐藏单元个数)的元素。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "37"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorShape([2, 256])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size = 2\n",
    "state = cell.get_initial_state(batch_size=batch_size,dtype=tf.float32)\n",
    "state.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "与上一节中实现的循环神经网络不同，这里`rnn_layer`的输入形状为(时间步数, 批量大小, 输入个数)。其中输入个数即one-hot向量长度（词典大小）。此外，`rnn_layer`作为Keras的`RNN`实例，在前向计算后会分别返回输出和隐藏状态，其中输出指的是隐藏层在各个时间步上计算并输出的隐藏状态，它们通常作为后续输出层的输入。需要强调的是，该“输出”本身并不涉及输出层计算，形状为(时间步数, 批量大小, 隐藏单元个数)。而`RNN`实例在前向计算返回的隐藏状态指的是隐藏层在最后时间步的可用于初始化下一时间步的隐藏状态：当隐藏层有多层时，每一层的隐藏状态都会记录在该变量中；对于像长短期记忆这样的循环神经网络，该变量还会包含其他信息。我们会在本章的后面介绍长短期记忆和深度循环神经网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "38"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([35, 2, 256]), 2, TensorShape([256]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_steps = 35\n",
    "X = tf.random.uniform(shape=(num_steps, batch_size, vocab_size))\n",
    "Y, state_new = rnn_layer(X, state)\n",
    "Y.shape, len(state_new), state_new[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来我们继承`Layer`类来定义一个完整的循环神经网络。它首先将输入数据使用one-hot向量表示后输入到`rnn_layer`中，然后使用全连接输出层得到输出。输出个数等于词典大小`vocab_size`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "39"
    }
   },
   "outputs": [],
   "source": [
    "# 本类已保存在d2lzh包中方便以后使用\n",
    "class RNNModel(keras.layers.Layer):\n",
    "    def __init__(self, rnn_layer, vocab_size, **kwargs):\n",
    "        super(RNNModel, self).__init__(**kwargs)\n",
    "        self.rnn = rnn_layer\n",
    "        self.vocab_size = vocab_size\n",
    "        self.dense = keras.layers.Dense(vocab_size)\n",
    "\n",
    "    def call(self, inputs, state):\n",
    "        # 将输入转置成(num_steps, batch_size)后获取one-hot向量表示\n",
    "        X = tf.one_hot(tf.transpose(inputs), self.vocab_size)\n",
    "        Y,state = self.rnn(X, state)\n",
    "        # 全连接层会首先将Y的形状变成(num_steps * batch_size, num_hiddens)，它的输出\n",
    "        # 形状为(num_steps * batch_size, vocab_size)\n",
    "        output = self.dense(tf.reshape(Y,(-1, Y.shape[-1])))\n",
    "        return output, state\n",
    "\n",
    "    def get_initial_state(self, *args, **kwargs):\n",
    "        return self.rnn.cell.get_initial_state(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练模型\n",
    "\n",
    "同上一节一样，下面定义一个预测函数。这里的实现区别在于前向计算和初始化隐藏状态的函数接口。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "41"
    }
   },
   "outputs": [],
   "source": [
    "# 本函数已保存在d2lzh包中方便以后使用\n",
    "def predict_rnn_keras(prefix, num_chars, model, vocab_size, idx_to_char,\n",
    "                      char_to_idx):\n",
    "    # 使用model的成员函数来初始化隐藏状态\n",
    "    state = model.get_initial_state(batch_size=1,dtype=tf.float32)\n",
    "    output = [char_to_idx[prefix[0]]]\n",
    "    #print(\"output:\",output)\n",
    "    for t in range(num_chars + len(prefix) - 1):\n",
    "        X = np.array([output[-1]]).reshape((1, 1))\n",
    "        #print(\"X\",X)\n",
    "        Y, state = model(X, state)  # 前向计算不需要传入模型参数\n",
    "        #print(\"Y\",Y)\n",
    "        #print(\"state:\",state)\n",
    "        if t < len(prefix) - 1:\n",
    "            output.append(char_to_idx[prefix[t + 1]])\n",
    "            #print(char_to_idx[prefix[t + 1]])\n",
    "        else:\n",
    "            output.append(int(np.array(tf.argmax(Y,axis=-1))))\n",
    "            #print(int(np.array(tf.argmax(Y[0],axis=-1))))\n",
    "    return ''.join([idx_to_char[i] for i in output])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "让我们使用权重为随机值的模型来预测一次。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "42"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'分开背伯欣苏人纳歉瞎屈钟'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = RNNModel(rnn_layer, vocab_size)\n",
    "predict_rnn_keras('分开', 10, model, vocab_size,  idx_to_char, char_to_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来实现训练函数。算法同上一节的一样，但这里只使用了相邻采样来读取数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 本函数已保存在d2lzh包中方便以后使用\n",
    "# 计算裁剪后的梯度\n",
    "def grad_clipping(grads,theta):\n",
    "    norm = np.array([0])\n",
    "    for i in range(len(grads)):\n",
    "        norm+=tf.math.reduce_sum(grads[i] ** 2)\n",
    "    #print(\"norm\",norm)\n",
    "    norm = np.sqrt(norm).item()\n",
    "    new_gradient=[]\n",
    "    if norm > theta:\n",
    "        for grad in grads:\n",
    "            new_gradient.append(grad * theta / norm)\n",
    "    else:\n",
    "        for grad in grads:\n",
    "            new_gradient.append(grad)  \n",
    "    #print(\"new_gradient\",new_gradient)\n",
    "    return new_gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "18"
    }
   },
   "outputs": [],
   "source": [
    "# 本函数已保存在d2lzh包中方便以后使用\n",
    "def train_and_predict_rnn_keras(model, num_hiddens, vocab_size, \n",
    "                                corpus_indices, idx_to_char, char_to_idx,\n",
    "                                num_epochs, num_steps, lr, clipping_theta,\n",
    "                                batch_size, pred_period, pred_len, prefixes):\n",
    "    loss = tf.keras.losses.SparseCategoricalCrossentropy()\n",
    "    optimizer=tf.keras.optimizers.SGD(learning_rate=lr)\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        l_sum, n, start = 0.0, 0, time.time()\n",
    "        data_iter = d2l.data_iter_consecutive(\n",
    "            corpus_indices, batch_size, num_steps)\n",
    "        state = model.get_initial_state(batch_size=batch_size,dtype=tf.float32)\n",
    "        for X, Y in data_iter:\n",
    "            with tf.GradientTape(persistent=True) as tape:\n",
    "                (outputs, state) = model(X, state)\n",
    "                y = Y.T.reshape((-1,))\n",
    "                l = loss(y,outputs)\n",
    "            \n",
    "            grads = tape.gradient(l, model.variables)\n",
    "            # 梯度裁剪\n",
    "            grads=grad_clipping(grads, clipping_theta)\n",
    "            optimizer.apply_gradients(zip(grads, model.variables))  # 因为已经误差取过均值，梯度不用再做平均\n",
    "            l_sum += np.array(l).item() * len(y)\n",
    "            n += len(y)\n",
    "\n",
    "        if (epoch + 1) % pred_period == 0:\n",
    "            print('epoch %d, perplexity %f, time %.2f sec' % (\n",
    "                epoch + 1, math.exp(l_sum / n), time.time() - start))\n",
    "            for prefix in prefixes:\n",
    "                print(' -', predict_rnn_keras(\n",
    "                    prefix, pred_len, model, vocab_size,  idx_to_char,\n",
    "                    char_to_idx))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用和上一节实验中一样的超参数来训练模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "attributes": {
     "classes": [],
     "id": "",
     "n": "19"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 50, perplexity 8313.898238, time 1.39 sec\n",
      " - 分开                                                  \n",
      " - 不分开                                                  \n",
      "epoch 100, perplexity 3665257.276584, time 1.41 sec\n",
      " - 分开于书觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉\n",
      " - 不分开觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉觉\n",
      "epoch 150, perplexity 1911317.223907, time 1.38 sec\n",
      " - 分开我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我\n",
      " - 不分开我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我我\n",
      "epoch 200, perplexity 2834346.220244, time 1.36 sec\n",
      " - 分开 小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小\n",
      " - 不分开 小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小\n",
      "epoch 250, perplexity 2834346.220244, time 1.36 sec\n",
      " - 分开 小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小\n",
      " - 不分开 小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小小\n"
     ]
    }
   ],
   "source": [
    "num_epochs, batch_size, lr, clipping_theta = 250, 32, 1e2, 1e-2\n",
    "pred_period, pred_len, prefixes = 50, 50, ['分开', '不分开']\n",
    "train_and_predict_rnn_keras(model, num_hiddens, vocab_size, \n",
    "                            corpus_indices, idx_to_char, char_to_idx,\n",
    "                            num_epochs, num_steps, lr, clipping_theta,\n",
    "                            batch_size, pred_period, pred_len, prefixes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 小结\n",
    "\n",
    "* Gluon的`rnn`模块提供了循环神经网络层的实现。\n",
    "* Gluon的`rnn.RNN`实例在前向计算后会分别返回输出和隐藏状态。该前向计算并不涉及输出层计算。\n",
    "\n",
    "## 练习\n",
    "\n",
    "* 与上一节的实现进行比较。看看Gluon的实现是不是运行速度更快？如果你觉得差别明显，试着找找原因。\n",
    "\n",
    "\n",
    "## 扫码直达[讨论区](https://discuss.gluon.ai/t/topic/4089)\n",
    "\n",
    "![](../img/qr_rnn-gluon.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": true,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
